# --------- MODEL 2 : SUDO-RM-RF -----------
# derived from: https://github.com/etzinis/sudo_rm_rf
# year: 2020
# architecture: based on conv-tas-net
# input: waveform
# output: waveform
# real-time factor: 1.84
sys.path.append('/Users/joanna.luberadzka/Projects/SpeechEnhanceTest/sudo_rm_rf-master/')
import sudo_rm_rf.dnn
import sudo_rm_rf.dnn.models.improved_sudormrf as improved_sudormrf
# load one of the pre-trained model (reverberant & noisy)
model_path = 'sudo_rm_rf-master/pretrained_models/Improved_Sudormrf_U36_Bases4096_WHAMRexclmark.pt'
anechoic_separation_model = torch.load(model_path)
# initialize pytorch model class
model = improved_sudormrf.SuDORMRF(
out_channels=anechoic_separation_model.out_channels,
in_channels=anechoic_separation_model.in_channels,
num_blocks=anechoic_separation_model.num_blocks,
upsampling_depth=anechoic_separation_model.upsampling_depth,
enc_kernel_size=anechoic_separation_model.enc_kernel_size,
enc_num_basis=anechoic_separation_model.enc_num_basis,
num_sources=anechoic_separation_model.num_sources,)
# load weights
model.load_state_dict(anechoic_separation_model.state_dict())
# move model to specified device (here i want to test with cpu/compare with mps)
model.to(device)
# switch on evaluation mode
model.eval()
def denoise_sudormrf(audio_path,fs_model):
# Sudo-rm-rf separator,
# Function to denoise with a little bit of pre and post-processing.
# ------------------------------------------------------
# load audio
fs, input_array = wavfile.read(audio_path)
# resample input audio to match model's fs
input_array = sig.resample_poly(input_array, fs_model, fs)
# array to tensor, prepare dimensions
input_tensor=torch.tensor(input_array).unsqueeze(0)
# Normalize the audio input tensor
input_std = input_tensor.std(-1, keepdim=True)
input_mean = input_tensor.mean(-1, keepdim=True)
input_tensor = (input_tensor - input_mean) / (input_std + 1e-9)
# move input data to specified device (here i want to test with cpu/compare with mps)
input_tensor=input_tensor.to(device)
# perform denoising and measure time:
t = time.time()
output_tensor = model(input_tensor.unsqueeze(1))
elapsed = time.time() - t
# Rescale the input sources with the mixture mean and variance
output_tensor = (output_tensor * input_std) + input_mean
# Go back to numpy array
output_array= output_tensor.detach().cpu().numpy()
# compute and print realtime factor
print("REALTIME FACTOR= "+ str((elapsed)/(output_array.shape[2]/fs_model)))
return input_array, output_array
# in a loop: denoise each file, compute realtime factor, display audio
fs_model=16000 #because this model works on a lower sampling rate then the sound
for filename in os.listdir("audios/noisy/"):
audio_path="audios/noisy/"+filename
noisy, separated=denoise_sudormrf(audio_path,fs_model)
# the output of the model is the separated speech and noise, so take speech:
target_denoised=separated[:,0,:]
wavfile.write(f"audios/denoised/sudormrf/sudormrf_{filename}", fs_model, target_denoised.T)
print("DENOISED FILE: " + filename)
display(Audio(data=noisy,rate=fs_model))
display(Audio(data=target_denoised,rate=fs_model))
# remove variables to make sure they are not used in the following cells:
%reset_selective -f "model"
%reset_selective -f "target_denoised"
